# src/rl_env/star_env.py
import numpy as np, pandas as pd, healpy as hp
import gymnasium as gym
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
from gymnasium import spaces
from scipy.spatial import cKDTree
from metrics.healpix_chi2 import healpix_chi2
from metrics.ripleyK_sphere import ripley_L
from metrics.wahba_cond import wahba_cond
from metrics.min_sep import min_sep
# 可选：from src.metrics.spec_coverage import spec_js

class StarEnv(gym.Env):
    """
    参数化筛选：动作调整 m_quantile + HEALPix cell 权重；奖励由若干指标加权组合。
    """
    metadata = {"render_modes": []}

    def __init__(self, catalog: pd.DataFrame, config: dict):
        super().__init__()
        self.cfg = config
        self.cat = catalog.reset_index(drop=True)
        assert {'ra','dec','mag'}.issubset(self.cat.columns)

        self.nside = self.cfg['nside']
        self.npix  = 12 * self.nside * self.nside
        self.topk  = self.cfg['topk_per_cell']
        self.min_sep = np.deg2rad(self.cfg['min_sep_arcsec']/3600.0)
        self.step_scale = self.cfg['step_scale']
        self.bright_goal = self.cfg['bright_ratio_goal']
        self.w_reward = self.cfg['w']

        # 预计算像素编号
        self.pix = hp.ang2pix(self.nside, self.cat['ra'].values, self.cat['dec'].values, lonlat=True)

        # 动作：δm_q + 像素权重增量
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(1+self.npix,), dtype=np.float32)
        # 观测向量：若要扩展，就在这里拼接更多指标
        self.observation_space = spaces.Box(low=-5, high=5, shape=(6,), dtype=np.float32)

        self.m_q = self.cfg['init_m_quantile']
        self.w   = np.zeros(self.npix, dtype=np.float32)
        self.sel = None

    def _to_xyz(self, df):
        ra, dec = np.deg2rad(df['ra'].values), np.deg2rad(df['dec'].values)
        x = np.cos(ra)*np.cos(dec); y = np.sin(ra)*np.cos(dec); z = np.sin(dec)
        return np.vstack([x,y,z]).T

    def _select_with_rules(self):
        m_thr = float(self.cat['mag'].quantile(self.m_q))
        mask_m = self.cat['mag'].values <= m_thr
        # 每个像素取前 topk
        order = np.argsort(-self.cat['mag'].values)  # 亮星在前
        pix_ordered = self.pix[order]
        top_idx = []
        for p in np.unique(pix_ordered):
            p_idx = np.where(pix_ordered == p)[0][:self.topk]
            top_idx.append(order[p_idx])
        if len(top_idx) == 0:
            return self.cat.iloc[[]]
        sel = self.cat.iloc[np.concatenate(top_idx)]
        return sel

    def _obs_stats(self, sel):
        chi2 = healpix_chi2(sel, self.nside)
        L    = ripley_L(sel)
        cond = wahba_cond(sel)
        N    = len(sel)
        bright = (sel['mag'] <= sel['mag'].quantile(0.2)).mean() if N>0 else 0.0

        # 可选项（没有就返回 0）
        sigma_m = 0.0   # 光度误差的标准差（占位）
        dmag_rmse = 0.0 # 仪器星等残差（占位）
        specjs = 0.0    # 谱型覆盖 JS divergence（占位）

        return dict(chi2=chi2, L=L, cond=cond, N=N, bright=bright,
                    sigma_m=sigma_m, dmag_rmse=dmag_rmse, specjs=specjs)

    def _reward(self, st):
        w = self.w_reward
        R = 0.0
        R += -w['chi2']   * st['chi2']
        R += -w['ripleyL']* st['L']
        R += -w['condK']  * np.log(st['cond']+1e-9)
        R +=  w['bright_goal'] * (st['bright'] >= self.bright_goal)
        R += -w['size']   * st['N']
        if w['phot']>0:   R +=  w['phot']  * np.exp(-(st['sigma_m']**2)/(0.05**2))
        if w['dmag']>0:   R += -w['dmag']  * st['dmag_rmse']
        if w['specjs']>0: R += -w['specjs']* st['specjs']
        return float(R)

    def _obs(self, st):
        v = np.array([
            st['chi2']*0.01, st['L']*10.0, np.log(st['cond']+1e-9)*0.1,
            st['N']/1e4, st['bright'], self.m_q
        ], dtype=np.float32)
        return np.tanh(v)

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.m_q = float(np.clip(self.m_q, 0.05, 0.95))
        self.w[:] = 0.0
        self.sel  = self._select_with_rules()
        st = self._obs_stats(self.sel)
        return self._obs(st), {}

    def step(self, action):
        delta_mq = float(action[0]) * self.step_scale
        self.m_q = float(np.clip(self.m_q + delta_mq, 0.05, 0.95))
        self.w  += np.asarray(action[1:], dtype=np.float32) * 0.1

        self.sel  = self._select_with_rules()
        st = self._obs_stats(self.sel)
        R  = self._reward(st)
        obs = self._obs(st)
        info = {"N": int(st['N']), "m_q": self.m_q, **st}
        return obs, R, False, False, info